# author：曹一鸣
import wmi
import json
import winreg
import datetime
import subprocess
import re
import pythoncom

class SystemInfo:
    def __init__(self):
        # 因为在线程里面没有初始化pythoncom，需要在你的线程任务里面再加一段代码，否则引入wmi后容易报错
        pythoncom.CoInitialize()
        self.__my_wmi = wmi.WMI()

    # 1. 获取计算机基本信息
    def get_base_info(self):
        wmi_win32_operating_system = self.__my_wmi.Win32_OperatingSystem()[0]
        # 获取计算机CPU数量,内存大小,主板相关信息
        wmi_win32_computer_system = self.__my_wmi.Win32_ComputerSystem()[0]
        # 用于获取CPU详细信息
        wmi_win32_processor = self.__my_wmi.Win32_Processor()[0]
        # 组装数据
        info = {
            # 操作系统版本
            "caption": wmi_win32_operating_system.Caption,
            # 计算机名字
            "CSName": wmi_win32_operating_system.CSName,
            # 获取操作系统类型(32bit/64bit)
            "OSArchitecture": wmi_win32_operating_system.OSArchitecture,
            # 获取操作系统语言版本
            "MUILanguages": wmi_win32_operating_system.MUILanguages[0],
            # 获取操作系统序列号
            "serialNumber": wmi_win32_operating_system.SerialNumber,
            # 获取主板厂商信息
            "manufacturer": wmi_win32_computer_system.Manufacturer,
            # 获取主板型号
            "model": wmi_win32_computer_system.Model,
            # 获取主板架构类型
            "systemType": wmi_win32_computer_system.SystemType,
            # 获取内存容量
            "totalPhysicalMemory": int(wmi_win32_computer_system.TotalPhysicalMemory) / 1024 / 1024,
            "cpuName": wmi_win32_processor.Name,  # 获取cpu类型
        }
        # ensure_ascii为False，确保dump Json字符串时A中文不被转换
        return json.dumps(info, ensure_ascii=False)

    # 2. 获取网络配置信息
    def get_network_info(self):
        #  配置及获取网络连接相关信息
        wmi_win32_network_adapter_configuration = self.__my_wmi.Win32_NetworkAdapterConfiguration(IPEnabled=1)
        infoList = []
        for item in wmi_win32_network_adapter_configuration:
            info = {
                "ipAddress": item.IPAddress[0],
                "netMask": item.IPSubnet[0],
                "mac": item.MACAddress
            }
            infoList.append(info)

        return json.dumps(infoList,ensure_ascii=False)

    # 3. 扫描已经安装软件的信息
    @staticmethod
    def get_install_soft():
        # 定义注册表安装、卸载软件检测位置
        sub_key = r'SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall'
        # 获取句柄handle(把注册表对象转换为python对象)
        key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
                             sub_key, 0, winreg.KEY_ALL_ACCESS)
        software_list = []
        # 通过winreg.QueryInfoKey(key)[0]获取当前注册表中的总个数
        for j in range(0, winreg.QueryInfoKey(key)[0] - 1):
            try:
                # 获取应用软件的尾缀名称
                key_name = winreg.EnumKey(key, j)
                # 组装完整的注册表软件路径
                key_path = sub_key + '\\' + key_name
                each_key = winreg.OpenKey(
                    winreg.HKEY_LOCAL_MACHINE, key_path, 0, winreg.KEY_ALL_ACCESS)
                display_name, rgz = winreg.QueryValueEx(
                    each_key, 'DisplayName')
                software_list.append(display_name)
            except WindowsError:
                pass

        keys = [str(x+1) for x in range(len(software_list))]
        software_list_json = dict(zip(keys, software_list))

        return json.dumps(software_list_json, ensure_ascii=False)

    # 4. 获取补丁信息
    @staticmethod
    def get_patch_info():
        patchInfo  = wmi.WMI().Win32_QuickFixEngineering()
        patchLists = []
        for patch in patchInfo:
            info = {
                "HotFixID":patch.HotFixID,
                "InstallDate":patch.InstalledOn,
                "Description":patch.Description,
                "InstalledBy":patch.InstalledBy
            }
            patchLists.append(info)

        return json.dumps(patchLists, ensure_ascii=False)

    # 5. 获取自启动项信息
    @staticmethod
    def get_run_once():
        # 5.1 连接注册表根键
        root1 = winreg.ConnectRegistry(None,winreg.HKEY_LOCAL_MACHINE)
        root2 = winreg.ConnectRegistry(None,winreg.HKEY_CURRENT_USER)
        result = {}
        try:
            # 5.2 指定想要访问的子健
            reg_path = r"SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run"
            reg_path1 = r"SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\RunOnce"
            reg_path2 = r"SOFTWARE\\Wow6432Node\\Microsoft\\Windows\\CurrentVersion\\Run"
            reg_path3 = r"SOFTWARE\\Wow6432Node\\Microsoft\\Windows\\CurrentVersion\\RunOnce"

            # 5.3 打开相应子健
            # 打开local machine的autorun列表
            key1 = winreg.OpenKey(root1, reg_path)
            # 打开current user的autorun列表
            key2 = winreg.OpenKey(root2, reg_path)
            key3 = winreg.OpenKey(root1, reg_path1)
            key4 = winreg.OpenKey(root2, reg_path1)
            key5 = winreg.OpenKey(root1, reg_path2)
            key6 = winreg.OpenKey(root1, reg_path3)
            key_list = [key1,key2,key3,key4,key5,key6]

            try:
                for key in key_list:
                    count = 0
                    while(1):
                        try:
                            # 5.4 通过winreg对象的EnumValue()方法迭代获取其中的键值
                            n, v, t = winreg.EnumValue(key, count)
                            result[n] = v
                            count += 1
                        except EnvironmentError:
                            break
            # finally是无论是否有异常，最后都要做的一些事情
            finally:
                # 5.5 关闭相应子健
                winreg.CloseKey(key1)
                winreg.CloseKey(key2)
                winreg.CloseKey(key3)
                winreg.CloseKey(key4)
                winreg.CloseKey(key5)
                winreg.CloseKey(key6)
        # finally是无论是否有异常，最后都要做的一些事情
        finally:
            # 5.6 关闭相应根键连接
            winreg.CloseKey(root1)
            winreg.CloseKey(root2)

        return json.dumps(result, ensure_ascii=False)

    # 6. 获取当前主机开放的端口
    @staticmethod
    def get_port_info():
        # 查看开放的、等待连接、但还没有被连接的端口
        cmd = 'netstat -ano | findstr LISTENING'
        res = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        cmd_str = [line.decode('utf-8') for line in res.communicate()][0]
        # 正则匹配得到含有端口号的字符串
        pattern = r"[0-9]+.[0-9]+.[0-9]+.[0-9]+:(\d+)"
        results = re.findall(pattern, cmd_str)
        results = list(set(results))
        results.remove("0")
        results = [int(res) for res in results]
        sorted_res = sorted(results)
        sorted_res = {str(i+1): sorted_res[i] for i in range(len(sorted_res))}
        return json.dumps(sorted_res, ensure_ascii=False)

    # 包含在系统基本信息获取：（1.） 获取系统唯一标识符UUID
    @staticmethod
    def get_uuid():
        cmd = 'wmic path win32_computersystemproduct get uuid'
        res = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        cmd_str = [line.decode('utf-8') for line in res.communicate()][0].splitlines()[2]
        return cmd_str


def main():
    base_info_scanner = SystemInfo()
    # 1. 获取计算机基本信息
    print(base_info_scanner.get_install_soft())
    # 2. 获取网络配置信息
    print(base_info_scanner.get_base_info())
    # 3. 扫描已经安装软件的信息
    print(base_info_scanner.get_network_info())
    # 4. 获取补丁信息
    print(base_info_scanner.get_patch_info())
    # 5. 获取自启动项信息
    print(base_info_scanner.get_run_once())
    # 6. 获取当前主机开放的端口
    print(base_info_scanner.get_port_info())



if __name__ == '__main__':
    main()
